import paddle
import paddle.nn as nn
import paddle.nn.functional as F



def focal_loss(inputs, targets, alpha=0.25, gamma=2.0, normalizer=None, use_sigmoid=False, reduction='mean'):
    if use_sigmoid:
        p = F.sigmoid(inputs)
    else:
        p = inputs
    ce_loss = F.binary_cross_entropy(p, targets, reduction='none')
    pt = p * targets + (1 - p) * (1 - targets)
    loss = ce_loss * (1 - pt)**gamma
    if alpha >= 0:
        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
        loss = alpha_t * loss
    
    if normalizer is None:
        normalizer = loss.numel()
    
    if reduction == 'mean':
        loss = loss.sum() / normalizer
    elif reduction == 'sum':
        loss = loss.sum()

    return loss
    
    
class FocalLoss(nn.Layer):
    def __init__(self, alpha=0.25, gamma=2.0, use_sigmoid=False, use_normalizer=True, reduction='mean'):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.use_sigmoid = use_sigmoid
        self.use_normalizer = use_normalizer
        self.reduction = reduction

    def forward(self, inputs, targets):
        if self.use_normalizer:
            normalizer = targets.sum()
        else:
            normalizer = None

        return focal_loss(inputs, targets, self.alpha, self.gamma, 
            normalizer=normalizer, use_sigmoid=self.use_sigmoid, reduction=self.reduction)

    
